import argparse
import platform
import torch
import yaml

from easydict import EasyDict as edict
from configparser import ConfigParser

from detector import YOLODetector

from . import PoseDetector


def update_config(config_file):
    with open(config_file) as f:
        config = edict(yaml.load(f, Loader=yaml.FullLoader))
        return config


def get_detector(opt=None):
    cfg_parser = ConfigParser()
    cfg_parser.read(r'yolo.cfg')
    try:
        yolo_cfg = cfg_parser['YOLO_CONFIG']
        tracker_cfg = cfg_parser['TRAKER_CONFIG']
    except Exception as err:
        print(err)
    else:
        # 缩减工程体量，暂不支持’traker‘模式
        if opt.detector == 'yolo':
            return YOLODetector(yolo_cfg, opt)
        else:
            # if opt.detector == 'tracker':
            #     return Tracker(tracker_cfg, opt)
            raise NotImplementedError


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='AlphaPose Demo')
    parser.add_argument('--cfg', type=str, required=True,
                        default='configs/halpe_136/resnet/256x192_res50_lr1e-3_2x-regression.yaml',
                        help='experiment configure file name')
    parser.add_argument('--gpus', type=str, dest='gpus', default="0",
                        help='choose which cuda device to use by index and input comma to use multi gpus, e.g. 0,1,2,3. (input -1 for cpu only)')
    parser.add_argument('--det_batch', type=int, default=5, help='detection batch size PER GPU')
    parser.add_argument('--pose_batch', type=int, default=64, help='pose estimation maximum batch size PER GPU')
    # 默认的识别器是yolo，暂不接入traker模式
    parser.add_argument('--detector', dest='detector', help='detector name', default="yolo")
    # parser.add_argument('--pose_flow', dest='pose_flow', help='track humans in video with PoseFlow',
    #                     action='store_true', default=False)
    # parser.add_argument('--pose_track', dest='pose_track', help='track humans in video with reid',
    #                     action='store_true', default=False)
    """
    设置模型运行参数
    """
    args = parser.parse_args()
    cfg = update_config(args.cfg)
    # 查询操作系统
    if platform.system() == 'Windows':
        args.sp = True
    # 查询显卡设备
    args.gpus = [int(i) for i in str(args.gpus).split(',')] if torch.cuda.device_count() >= 1 else [-1]
    # 设置计算设备（cpu/gpu）
    args.device = torch.device("cuda:" + str(args.gpus[0]) if args.gpus[0] >= 0 else "cpu")
    # 设置批处理规模
    args.det_batch = args.detbatch * len(args.gpus)
    args.pose_batch = args.posebatch * len(args.gpus)

    """
    设置计算单元共享策略
    """
    if not args.sp:
        torch.multiprocessing.set_start_method('forkserver', force=True)
        torch.multiprocessing.set_sharing_strategy('file_system')

    """
    加载模型
    """
    pose_detector = get_detector(args)
    # 获取待识别图像
    input_source = list([])
    # 只保留输入数据为image的识别方法，其他格式（frame、video）暂不实现
    det_loader = PoseDetector(pose_detector,
                              args=args,
                              cfg=cfg,
                              batchSize=args.detbatch,
                              queueSize=args.qsize,
                              mode='image'
                              )
    det_worker = det_loader.start(frames=input_source)
